# tensor_visualizer.py

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from mpl_toolkits.mplot3d import Axes3D
from scipy.stats import skew, kurtosis
import warnings

# --- Global Aesthetic Settings ---
# Set Seaborn style for better visual aesthetics
# Common styles: "darkgrid", "whitegrid", "dark", "white", "ticks"
sns.set_theme(style="whitegrid", palette="viridis")
# plt.rcParams["font.sans-serif"] = ["SimHei"]  # For displaying Chinese labels normally
# plt.rcParams["axes.unicode_minus"] = False  # For displaying negative signs normally
warnings.filterwarnings("ignore", category=UserWarning, module="matplotlib")


def get_outliers_iqr(data: np.ndarray, threshold: float = 1.5):
    """
    Identify outliers using the IQR (Interquartile Range) method.

    Args:
        data (np.ndarray): Input 1D numpy array.
        threshold (float): IQR multiplier threshold, typically 1.5 (standard) or 3.0 (more extreme).

    Returns:
        tuple: (outliers, outlier_indices)
    """
    q1, q3 = np.percentile(data, [25, 75])
    iqr = q3 - q1
    lower_bound = q1 - threshold * iqr
    upper_bound = q3 + threshold * iqr

    outlier_indices = np.where((data < lower_bound) | (data > upper_bound))[0]
    outliers = data[outlier_indices]

    return outliers, outlier_indices


def print_tensor_stats(
    tensor: torch.Tensor, name: str = "Tensor", outlier_threshold: float = 3.0
):
    """
    Print detailed statistical information about a tensor.

    Args:
        tensor (torch.Tensor): PyTorch tensor to analyze.
        name (str): Tensor name for printing.
        outlier_threshold (float): IQR multiplier for defining outliers.
    """
    if not isinstance(tensor, torch.Tensor):
        raise TypeError("Input must be a torch.Tensor")

    # Move tensor to CPU and convert to numpy array for analysis
    data = tensor.detach().cpu().numpy()
    flat_data = data.flatten()

    print("\n" + "=" * 30)
    print(f" Statistical Analysis of '{name}'")
    print("=" * 30)
    print(f" - Shape: {data.shape}")
    print(f" - dtype: {tensor.dtype}")
    print(f" - Device: {tensor.device}")
    print(f" - Total elements: {data.size}")

    print("\n--- Numerical Statistics ---")
    print(f" - Min: {np.min(data):.4f}")
    print(f" - Max: {np.max(data):.4f}")
    print(f" - Mean: {np.mean(data):.4f}")
    print(f" - Median: {np.median(data):.4f}")
    print(f" - Std Dev: {np.std(data):.4f}")

    print("\n--- Quantiles ---")
    quantiles = np.percentile(flat_data, [0, 1, 25, 50, 75, 99, 100])
    print(f" - 0% (Min): {quantiles[0]:.4f}")
    print(f" - 1%      : {quantiles[1]:.4f}")
    print(f" - 25% (Q1): {quantiles[2]:.4f}")
    print(f" - 50% (Med):{quantiles[3]:.4f}")
    print(f" - 75% (Q3): {quantiles[4]:.4f}")
    print(f" - 99%     : {quantiles[5]:.4f}")
    print(f" - 100% (Max):{quantiles[6]:.4f}")

    print("\n--- Distribution Shape ---")
    print(f" - Skewness: {skew(flat_data):.4f}")
    print(f" - Kurtosis: {kurtosis(flat_data):.4f}")

    print("\n--- Special Values & Sparsity ---")
    print(f" - NaN count: {np.isnan(data).sum()}")
    print(f" - Inf count: {np.isinf(data).sum()}")
    sparsity = 100.0 * (np.count_nonzero(data == 0) / data.size)
    print(f" - Sparsity: {sparsity:.2f}%")

    print("\n--- Outlier Analysis (IQR > 1.5) ---")
    outliers, outlier_indices = get_outliers_iqr(flat_data, threshold=1.5)
    print(f" - Outlier count (IQR > 1.5): {len(outliers)}")
    if len(outliers) > 0:
        print(f" - Outlier percentage: {100 * len(outliers) / len(flat_data):.4f}%")
        print(f" - Outlier range: [{np.min(outliers):.4f}, {np.max(outliers):.4f}]")

    print("=" * 30 + "\n")


def plot_tensor_distribution(
    tensor: torch.Tensor, name: str = "Tensor", bins: int = 100
):
    """
    Plot histogram, boxplot, and violin plot of tensor values.

    Args:
        tensor (torch.Tensor): PyTorch tensor to visualize.
        name (str): Tensor name for plot titles.
        bins (int): Number of bins for histogram.
    """
    if not isinstance(tensor, torch.Tensor):
        raise TypeError("Input must be a torch.Tensor")

    data = tensor.detach().cpu().flatten().numpy()

    fig, axes = plt.subplots(1, 3, figsize=(20, 6))
    fig.suptitle(f"'{name}' - Data Distribution Visualization", fontsize=16)

    # 1. Histogram
    sns.histplot(data, bins=bins, kde=True, ax=axes[0])
    axes[0].set_title("Histogram with KDE")
    axes[0].set_xlabel("Value")
    axes[0].set_ylabel("Frequency")
    axes[0].axvline(
        data.mean(), color="r", linestyle="--", label=f"Mean: {data.mean():.2f}"
    )
    axes[0].axvline(
        np.median(data),
        color="g",
        linestyle="-",
        label=f"Median: {np.median(data):.2f}",
    )
    axes[0].legend()

    # 2. Boxplot
    sns.boxplot(x=data, ax=axes[1], orient="h")
    axes[1].set_title("Boxplot (Highlighting Outliers)")
    axes[1].set_xlabel("Value")

    # 3. Violin plot
    sns.violinplot(x=data, ax=axes[2], orient="h", inner="quartile")
    axes[2].set_title("Violin Plot")
    axes[2].set_xlabel("Value")

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    plt.show()


def plot_scatter_with_outliers(
    tensor: torch.Tensor,
    name: str = "Tensor",
    outlier_threshold: float = 1.5,
    annotate_top_k: int = 5,
):
    """
    Plot a 2D scatter plot of tensor values with highlighted and annotated outliers.

    Args:
        tensor (torch.Tensor): PyTorch tensor to visualize.
        name (str): Tensor name for plot title.
        outlier_threshold (float): IQR multiplier for defining outliers.
        annotate_top_k (int): Number of most extreme outliers to annotate.
    """
    if not isinstance(tensor, torch.Tensor):
        raise TypeError("Input must be a torch.Tensor")

    original_shape = tensor.shape
    data = tensor.detach().cpu().numpy()
    flat_data = data.flatten()

    outliers, outlier_indices = get_outliers_iqr(flat_data, threshold=outlier_threshold)

    plt.figure(figsize=(14, 8))

    # Plot all points
    all_indices = np.arange(len(flat_data))
    plt.scatter(all_indices, flat_data, alpha=0.5, label="Normal values", s=15)

    if len(outliers) > 0:
        # Highlight outliers
        plt.scatter(
            outlier_indices,
            outliers,
            color="red",
            s=40,
            label=f"Outliers (IQR > {outlier_threshold})",
            zorder=5,
        )

        # Annotate most extreme outliers
        if annotate_top_k > 0:
            # Find top k outliers by absolute value for annotation
            abs_outliers = np.abs(outliers)
            top_k_indices_in_outliers = np.argsort(abs_outliers)[-annotate_top_k:]

            for idx in top_k_indices_in_outliers:
                original_outlier_index = outlier_indices[idx]
                outlier_value = flat_data[original_outlier_index]
                # Convert flat index back to multi-dimensional coordinates
                multi_dim_coords = np.unravel_index(
                    original_outlier_index, original_shape
                )

                plt.annotate(
                    f"Coords: {multi_dim_coords}\nValue: {outlier_value:.2f}",
                    xy=(original_outlier_index, outlier_value),
                    xytext=(10, 10),
                    textcoords="offset points",
                    arrowprops=dict(arrowstyle="->", connectionstyle="arc3,rad=.2"),
                    bbox=dict(
                        boxstyle="round,pad=0.3",
                        fc="yellow",
                        ec="black",
                        lw=1,
                        alpha=0.7,
                    ),
                )

    plt.title(f"'{name}' - Element Scatter Plot with Outlier Annotation", fontsize=16)
    plt.xlabel("Element Index (flattened)")
    plt.ylabel("Value")
    plt.legend()
    plt.grid(True)
    plt.show()


def plot_3d_surface(tensor: torch.Tensor, name: str = "Tensor"):
    """
    Plot a 3D surface plot for 2D tensors.

    Args:
        tensor (torch.Tensor): 2D PyTorch tensor to visualize.
        name (str): Tensor name for plot title.
    """
    if not isinstance(tensor, torch.Tensor):
        raise TypeError("Input must be a torch.Tensor")
    if tensor.dim() != 2:
        print(
            f"Note: 3D surface plot only works for 2D tensors, but '{name}' has {tensor.dim()} dimensions. Skipping this plot."
        )
        return

    data = tensor.detach().cpu().numpy()
    x = np.arange(data.shape[1])
    y = np.arange(data.shape[0])
    X, Y = np.meshgrid(x, y)

    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection="3d")

    surf = ax.plot_surface(X, Y, data, cmap="viridis", edgecolor="none")

    ax.set_title(f"'{name}' - 3D Surface Plot", fontsize=16)
    ax.set_xlabel("Dimension 1 Index")
    ax.set_ylabel("Dimension 0 Index")
    ax.set_zlabel("Value")
    fig.colorbar(surf, shrink=0.5, aspect=5, label="Value")

    plt.show()


if __name__ == "__main__":
    """
    Example usage of the above functions.
    We create a simulated tensor with obvious outliers.
    """
    # --- Example 1: A near-normal distribution weight tensor with some extreme outliers ---
    print("*" * 15 + " Example 1: Simulated Weight Tensor " + "*" * 15)
    # Create a base tensor
    weights = torch.randn(128, 64) * 0.5  # Simulate most weights

    # Inject some extreme outliers
    num_outliers = 10
    for _ in range(num_outliers):
        idx1 = torch.randint(0, 128, (1,)).item()
        idx2 = torch.randint(0, 64, (1,)).item()
        # Random positive/negative outliers
        sign = 1 if torch.rand(1).item() > 0.5 else -1
        weights[idx1, idx2] = torch.randn(1) * 5 + sign * 10

    # 1. Print statistics
    print_tensor_stats(weights, name="Simulated Weights")

    # 2. Plot distribution
    plot_tensor_distribution(weights, name="Simulated Weights", bins=150)

    # 3. Plot scatter with outlier annotations
    plot_scatter_with_outliers(weights, name="Simulated Weights", annotate_top_k=5)

    # 4. Plot 3D surface (since it's 2D)
    plot_3d_surface(weights, name="Simulated Weights")

    # --- Example 2: A post-ReLU activation feature map tensor (non-negative, sparse) ---
    print("\n" + "*" * 15 + " Example 2: Simulated ReLU Activation Tensor " + "*" * 15)
    activations = (
        torch.randn(32, 64, 16, 16) - 0.2
    )  # Shift center left to produce more zeros
    activations = torch.nn.functional.relu(activations)  # ReLU activation

    # Inject some large activation values
    activations[5, 10, 8, 8] = 25.0
    activations[15, 20, 12, 12] = 30.0

    # 1. Print statistics
    print_tensor_stats(activations, name="Simulated Activations")

    # 2. Plot distribution
    plot_tensor_distribution(activations, name="Simulated Activations")

    # 3. Plot scatter
    plot_scatter_with_outliers(
        activations, name="Simulated Activations", annotate_top_k=3
    )

    # 4. Plot 3D surface (will skip since it's 4D)
    plot_3d_surface(activations, name="Simulated Activations")
